package bayesian;

import java.util.ArrayList;
import java.util.Random;

import probability.GaussianDistribution;
import probability.ProbabilityDistribution;
import utils.Printer;
import basics.DenseMatrix;
import basics.DenseVector;
import basics.FeaturesSet;
import basics.Matrix;
import basics.Vector;
import basics.VectorMatrix;
import basics.VectorMatrixUtils.OperationOpt;

public class SimpleBayesianSampling {

	/**
	 * This method computes the updated value for the i,jth latent variable.
	 * 
	 * @param i
	 *            : index i
	 * @param j
	 *            : index j
	 * @param prior
	 *            : The prior probability for the parameter
	 * @param UV
	 *            : The matrix on which the operation has to be performed, either U or V
	 * @param xy
	 *            : The input vector, either x or y
	 * @param z
	 *            : The other side of the latent model > xUVy => z = Vy if V is frozen and we want to calculate U, and z
	 *            = xU otherwise.
	 */
	private static GaussianDistribution newUV(int i, int j, GaussianDistribution prior,
			final GaussianDistribution[][] UV, final Vector xy, final VectorMatrix z, int zIdx, double r, double sum) {
		GaussianDistribution p = new GaussianDistribution(0, 1);

		double s = xy.get(i) * z.get(zIdx);
		if (s == 0)
			return UV[i][j];

		// s = s == 0 ? 1 : s;
		p.setMu(r / s);
		double n = 0;

		for (int k = 0; k < xy.length(); k++) {
			for (int l = 0; l < z.length(); l++) {
				if (k != i && l != j) {
					double t = xy.get(k) * z.get(l);
					if (xy.length() == UV.length) {
						t *= UV[k][l].getMu();
					} else {
						t *= UV[l][k].getMu();
					}
					if (t != 0) {
						n++;
						p.setMu(p.getMu() - t / s);
					}
				}
			}
		}

		double delta = p.getMu() - UV[i][j].getMu();
		p.setMu(Math.abs(UV[i][j].getMu()) * delta / sum + UV[i][j].getMu());
		p = GaussianDistribution.product(prior, p);

		return p;
	}

	GaussianDistribution[][] _U;
	DenseMatrix _UMean;
	GaussianDistribution[][] _V;
	DenseMatrix _VMean;

	public void train(FeaturesSet X, FeaturesSet Y, FeaturesSet R, int lDimension) {
		_U = initLatents(X.get(0).length(), lDimension);
		_V = initLatents(lDimension, Y.get(0).length());
		_UMean = getMeanMatrix(_U);
		_VMean = getMeanMatrix(_V);

		for (int i = 0; i < X.size(); i++) {
			for (int j = 0; j < Y.size(); j++) {
				if (R.get(i).get(j) != 0) {
					System.out.println(i + "," + j);
					update((Vector) X.get(i), (Vector) Y.get(j), R.get(i).get(j), lDimension);
					Printer.appendToFile("uv.txt", "Iteration " + i + ", " + j + toString(_U)
							+ "\n-------------------------------------");
					Printer.appendToFile("uv.txt", toString(_V) + "\n-------------------------------------");
				}
			}
		}

		_UMean = getMeanMatrix(_U);
		_VMean = getMeanMatrix(_V);

		_U = _V = null;

	}

	private void update(Vector x, Vector y, double r, int lDimension) {
		ArrayList<GaussianDistribution[][]> usamples = new ArrayList<GaussianDistribution[][]>();
		ArrayList<GaussianDistribution[][]> vsamples = new ArrayList<GaussianDistribution[][]>();

		for (int sample = 0; sample < 1000; sample++) {

			GaussianDistribution[][] u = new GaussianDistribution[x.length()][lDimension]; // _U; //
			GaussianDistribution[][] v = new GaussianDistribution[lDimension][y.length()]; // _V; //
			DenseMatrix umean = getMeanMatrix(_U);
			DenseMatrix vmean = getMeanMatrix(_V);
			double usum = umean.sumAllElements(OperationOpt.AbsVal);
			double vsum = vmean.sumAllElements(OperationOpt.AbsVal);

			for (int i = 0; i < x.length(); i++) {
				for (int k = 0; k < lDimension; k++) {
					Vector z = vmean.times(y.transpose()).getColumnVector();
					u[i][k] = newUV(i, k, u[i][k], u, x, z, k, r, usum); // new GaussianDistribution(1, 1)
					umean.set(u[i][k].getMu(), i, k);
				}
			}

			_U = u;
			usamples.add(copy(u));

			for (int j = 0; j < y.length(); j++) {
				for (int k = 0; k < lDimension; k++) {
					Vector z = x.times(umean).getColumnVector();
					v[k][j] = newUV(k, j, v[k][j], v, y, z, k, r, vsum); // new GaussianDistribution(1, 1)
					vmean.set(v[k][j].getMu(), k, j);
				}
			}
			_V = v;
			vsamples.add(copy(v));
		}

		computeFromSamples(_U, usamples);
		computeFromSamples(_V, vsamples);

		_UMean = getMeanMatrix(_U);
		_VMean = getMeanMatrix(_V);
	}

	private static void computeFromSamples(GaussianDistribution[][] UV, ArrayList<GaussianDistribution[][]> samples) {
		final int start = 100;
		for (int i = 0; i < UV.length; i++) {
			for (int j = 0; j < UV[0].length; j++) {
				double mean = 0, s = 0;
				for (int sample = start; sample < samples.size(); sample++) {
					double d = samples.get(sample)[i][j].getMu() - mean;
					mean += (d / (double) (sample - start + 1));
					s += (d * (samples.get(sample)[i][j].getMu() - mean) / (double) (sample - start + 1));
				}
				UV[i][j].setMu(mean);
				UV[i][j].setSigma2(s);
			}
		}
	}

	private GaussianDistribution[][] copy(GaussianDistribution[][] p) {
		GaussianDistribution[][] cp = new GaussianDistribution[p.length][p[0].length];

		for (int i = 0; i < p.length; i++) {
			for (int j = 0; j < p[0].length; j++) {
				cp[i][j] = (GaussianDistribution) p[i][j].clone();
			}
		}

		return cp;
	}

	private String toString(GaussianDistribution[][] p) {
		StringBuilder sb = new StringBuilder();
		for (int i = 0; i < p.length; i++) {
			for (int j = 0; j < p[0].length; j++) {
				sb.append(p[i][j].toString());
			}
		}
		return sb.toString();
	}

	private DenseMatrix getMeanMatrix(GaussianDistribution[][] p) {
		DenseMatrix m = new DenseMatrix(p.length, p[0].length);

		for (int i = 0; i < p.length; i++) {
			for (int j = 0; j < p[0].length; j++) {
				m.set(p[i][j].getMu(), i, j);
			}
		}

		return m;
	}

	public double predict(Vector x, Vector y) {
		double m = x.times(_UMean).times(_VMean).times(y.transpose()).get(0, 0);
		if (m > 5)
			m = 5;
		else if (m < 1)
			m = 1;
		return m;
	}

	private GaussianDistribution[][] initLatents(int n, int m) {
		GaussianDistribution[][] p = new GaussianDistribution[n][m];
		Random r = new Random();
		for (int i = 0; i < n; i++) {
			for (int j = 0; j < m; j++) {
				p[i][j] = new GaussianDistribution(r.nextDouble(), 10.);
			}
		}
		return p;
	}

	public static String trainEval(FeaturesSet X, FeaturesSet Y, FeaturesSet R, FeaturesSet X_eval, FeaturesSet Y_eval,
			FeaturesSet R_eval, int k) {

		SimpleBayesianSampling sb = new SimpleBayesianSampling();
		sb.train(X, Y, R, k);

		double mae = 0;
		double n = 0;
		for (int i = 0; i < X_eval.size(); i++) {
			for (int j = 0; j < Y_eval.size(); j++) {
				if (R_eval.get(i).get(j) <= 0)
					continue;
				double d = sb.predict((Vector) X_eval.get(i), (Vector) Y_eval.get(j));
				double err = Math.abs(d - R_eval.get(i).get(j));
				mae += err;
				n++;
			}
		}

		return "MAE: " + mae / n;
	}
}
